#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Copyright 2020 Huawei Technologies Co., Ltd
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import sys
import acl
import numpy as np
from common.const import PIXEL_FORMAT_YUV_SEMIPLANAR_420, \
                                     PIXEL_FORMAT_YVU_SEMIPLANAR_420, \
                                     PIXEL_FORMAT_YUV_SEMIPLANAR_422, \
                                     PIXEL_FORMAT_YVU_SEMIPLANAR_422, \
                                     PIXEL_FORMAT_YUV_SEMIPLANAR_444, \
                                     PIXEL_FORMAT_YVU_SEMIPLANAR_444, \
                                     STATIC_BATCH_RESOLUTION, DYNAMIC_BATCH, \
                                     DYNAMIC_RESOLUTION, \
                                     APPEND, REMOVE, \
                                     DEBUG, INFO, WARNING, ERROR, \
                                     DEVICE, HOST, DVPP, \
                                     ACL_MEM_MALLOC_HUGE_FIRST, \
                                     ACL_MEM_MALLOC_HUGE_ONLY, \
                                     ACL_MEM_MALLOC_NORMAL_ONLY
from common.log import Log, check_ret
sys.path.append("..")


class SimpleMemoryPool():


    def __init__(self, mode=STATIC_BATCH_RESOLUTION):
        self.mode = mode
        self.class_name = self.__class__.__name__

    def buffer_size_calculation(self, 
                                width_stride: int, 
                                height_stride: int,
                                data_format: int):
        if width_stride is None or height_stride is None or \
                width_stride <= 0 or height_stride <= 0:
            Log.do_log(WARNING, 
                       self.class_name, 
                       'buffer_size_calculation',
                       'width_stride = {} and height_stride = {}' \
                       .format(width_stride, height_stride))
            width_stride = 0
            height_stride = 0

        if data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_420 or \
            data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_420:
            buffer_size = width_stride * height_stride * 3 // 2
        elif data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_422 or \
            data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_422:
            buffer_size = width_stride * height_stride * 2
        elif data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_444 or \
            data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_444:
            buffer_size = width_stride * height_stride * 3
        else:
            Log.do_log(WARNING, self.class_name, 'buffer_size_calculation', \
                    'data_format is not supported')
            buffer_size = 0
        return buffer_size

    def stride_calculation(self, height, width, func): 
        if func == 'jpegd':
            width_alignment = 128
            height_alignment = 16
        elif func == 'vdec' or \
             func == 'resize' or \
             func == 'crop' or \
             func == 'crop_and_paste':
            width_alignment = 16
            height_alignment = 2
        else:
            Log.do_log(WARNING, 
                       self.class_name, 
                       'stride_calculation', 
                       'func is not supported')
            return 0, 0
        width_stride = ((width + width_alignment - 1) // width_alignment) * width_alignment
        height_stride = ((height + height_alignment - 1) // height_alignment) * height_alignment
        return width_stride, height_stride

    def memory_allocation(self, 
                          func: str, 
                          height: int, 
                          width: int, 
                          data_format: int):
        width_stride, height_stride = self.stride_calculation(height, width, func)
        buffer_size = self.buffer_size_calculation(width_stride, height_stride, data_format)
        start_ptr = self.memory_allocation_(buffer_size, DVPP)
        return start_ptr, buffer_size, width_stride, height_stride
    
    def memory_allocation_(self, size: int, place: int):
        if place == DVPP:
            ptr, ret = acl.media.dvpp_malloc(size)
            check_ret(self.class_name, 'memory_allocation', 'fail to malloc dvpp memory', ret)
        elif place == DEVICE:
            ptr, ret = acl.rt.malloc(size, ACL_MEM_MALLOC_HUGE_FIRST)
            check_ret(self.class_name, 'memory_allocation', 'fail to malloc device memory', ret)
        elif place == HOST:
            ptr, ret = acl.rt.malloc_host(size)
            check_ret(self.class_name, 'memory_allocation', 'fail to malloc host memory', ret)       
        else:
            check_ret(self.class_name, 'memory_allocation', 'out of options', 1)
        return ptr

    def memory_reset(self, ptr: int, size: int, value):
        if ptr and size:
            ret = acl.rt.memset(ptr, size, value, size)
            check_ret(self.class_name, 'memory_reset', 'fail to reset by acl.rt.memset', ret)
        else:
            Log.do_log(WARNING, self.class_name, 'memory_reset', 'ptr = {} size = {}'.format(ptr, size))

    def release_memory(self, buffer_ptr, place: int):
        if buffer_ptr != None:
            if place == DVPP:
                ret = acl.media.dvpp_free(buffer_ptr)
                check_ret(self.class_name, 'release_memory', 'fail to release dvpp memory', ret)
            elif place == DEVICE:
                ret = acl.rt.free(buffer_ptr)
                check_ret(self.class_name, 'release_memory', 'fail to release memory', ret)
            elif place == HOST:
                ret = acl.rt.free_host(buffer_ptr)
                check_ret(self.class_name, 'release_memory', 'fail to release memory', ret) 
            else:
                Log.do_log(WARNING, self.class_name, 'release_memory', 'out of option')

